from TLSMapper.TLSSUT import *
from TLSMapper.TLSProtocol import *
from datetime import datetime
from aalpy.utils import visualize_automaton
from aalpy.learning_algs import run_Lstar
from aalpy.oracles import RandomWMethodEqOracle, StatePrefixEqOracle
from aalpy.base import SUL
import os
import argparse
import subprocess
import signal
import time
import socket
import logging

from learning.StatePrefixEqOracleFailSafe import StatePrefixOracleFailSafe
from learning.FailSafeCacheSUL import FailSafeCacheSUL
from logger import init_global_logger


# 目标查询序列（遇到此序列时停止）
TARGET_QUERY = ('ClientHello', 'ClientKeyExchange', 'ChangeCipherSpec', 'ClosureAlert', 'Finish')


class TargetQueryReached(Exception):
    """自定义异常：当遇到目标查询序列时抛出"""
    def __init__(self, target_query, query_count, output=None):
        self.target_query = target_query
        self.query_count = query_count
        self.output = output
        super().__init__(f"遇到目标查询序列: {target_query} (查询数: {query_count})")


# 'OpenSSL', 'WolfSSL', 'MatrixSSL', 'GnuTLS', 'mbedTLS', 'BoringSSL', 'tlse'
ciphersuites=[
              CipherSuite.TLS_AES_128_GCM_SHA256,
              CipherSuite.TLS_AES_256_GCM_SHA384,
              CipherSuite.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
              CipherSuite.TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
              CipherSuite.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384
              ]

# 证书路径配置（写死）
CERT_BASE = '/home/ubuntu/experiments/certs'
CLIENT_KEY = f'{CERT_BASE}/client-key.pem'
CLIENT_CERT = f'{CERT_BASE}/client-cert.pem'
SERVER_CERT = f'{CERT_BASE}/server-cert.pem'
SERVER_KEY = f'{CERT_BASE}/server-key.pem'
CA_CERT = f'{CERT_BASE}/ca-cert.pem'
SERVER_HOST = '127.0.0.1'
SERVER_PORT = 1111

# 服务器路径配置
TLSE_ROOT = '/home/ubuntu/experiments/tlse'
SERVER_BIN = f'{TLSE_ROOT}/build/examples/server'

# TLS 1.2 字母表（不包含 TLS 1.3 专用消息）
TLS12_ALPHABET = [
    'ClientHello',
    'Certificate',
    'ClientKeyExchange',
    'CertificateVerify',
    'ChangeCipherSpec',
    'Finish',
    'ApplicationData',
    'ResumptionClientHello',
    'EmptyCertificate',
    'ClosureAlert'
]


class DebugStopSUL(SUL):
    """
    包装 SUL，用于在遇到目标查询序列时停止学习过程
    """
    def __init__(self, sul: SUL, target_query):
        super().__init__()
        self.sul = sul
        self.target_query = tuple(target_query) if isinstance(target_query, (list, tuple)) else target_query
        self.query_count = 0
        self.stopped = False
        
    def query(self, word):
        """
        执行查询，如果遇到目标序列则停止
        """
        self.query_count += 1
        
        # 将 word 转换为元组以便比较
        if isinstance(word, (list, tuple)):
            word_tuple = tuple(word)
        else:
            word_tuple = (word,)
        
        # 检查是否匹配目标查询
        if word_tuple == self.target_query:
            print("\n" + "=" * 80)
            print("⚠️  检测到目标查询序列！")
            print(f"目标序列: {self.target_query}")
            print(f"当前查询: {word_tuple}")
            print(f"已执行查询数: {self.query_count}")
            print("=" * 80)
            
            # 执行查询以获取响应
            print("\n执行目标查询以获取响应...")
            output = None
            try:
                output = self.sul.query(word)
                print(f"查询响应: {output}")
                print("=" * 80)
            except Exception as e:
                print(f"执行查询时出错: {e}")
                import traceback
                traceback.print_exc()
            
            # 标记已停止
            self.stopped = True
            
            # 抛出自定义异常以停止 L* 算法
            raise TargetQueryReached(self.target_query, self.query_count, output)
        
        # 正常执行查询
        return self.sul.query(word)
    
    def pre(self):
        """重置系统"""
        self.sul.pre()
    
    def post(self):
        """后处理"""
        self.sul.post()
    
    def step(self, letter):
        """执行单步"""
        return self.sul.step(letter)


def get_happy_flow():
    """返回 tls12_oneway 的 happy flow 序列"""
    return [
        'ClientHello',
        'ClientKeyExchange',
        'ChangeCipherSpec',
        'Finish',
        'ApplicationData'
    ]


def get_mode_config():
    """返回 tls12_oneway 模式的配置信息"""
    # 基础服务器命令（单向认证）
    base_server_cmd_oneway = [
        SERVER_BIN,
        '0.0.0.0',
        str(SERVER_PORT),
        SERVER_CERT,
        SERVER_KEY,
        '-f'
    ]
    
    return {
        'version': (3, 3),  # TLS 1.2
        'need_client_cert': False,
        'alphabet': TLS12_ALPHABET,
        'description': 'TLS 1.2 - 单向认证（服务器证书）',
        'server_cmd': base_server_cmd_oneway
    }


def create_tls_protocol(mode_config):
    """创建 TLS 协议配置"""
    return TLSProtocol(
        version=mode_config['version'],
        cipher_suite=ciphersuites,
        implementation='tlse',
        verify=False,
        fuzz_tls=(False, None, None),
        target=(SERVER_HOST, SERVER_PORT)
    )


def create_sut(tlsimp, mode_config):
    """创建 SUT（System Under Test）"""
    return TLSSUT(
        TLSpro=tlsimp,
        keyfile=None,
        certfile=None,
        server_cert=CA_CERT
    )


def wait_for_port(host, port, timeout=10):
    """等待端口可用"""
    start_time = time.time()
    attempt = 0
    while time.time() - start_time < timeout:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex((host, port))
            sock.close()
            if result == 0:
                elapsed = time.time() - start_time
                if attempt > 0:
                    print(f"  端口 {host}:{port} 已可用 (等待了 {elapsed:.2f} 秒)")
                return True
        except Exception as e:
            pass
        attempt += 1
        time.sleep(0.2)
    
    elapsed = time.time() - start_time
    print(f"  端口 {host}:{port} 在 {elapsed:.2f} 秒内未变为可用")
    return False


def start_server(server_cmd):
    """启动服务器进程"""
    print(f"启动服务器: {' '.join(server_cmd)}")
    proc = subprocess.Popen(
        server_cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        preexec_fn=os.setsid,
        close_fds=True
    )
    
    time.sleep(1.0)
    
    if proc.poll() is not None:
        print(f"✗ 服务器进程已退出 (退出码: {proc.returncode})")
        return None
    
    print(f"  等待服务器在 {SERVER_HOST}:{SERVER_PORT} 上监听...")
    if wait_for_port(SERVER_HOST, SERVER_PORT, timeout=15):
        print(f"✓ 服务器已启动 (PID: {proc.pid})")
        return proc
    else:
        print("✗ 服务器启动超时")
        print(f"  检查进程状态: {'运行中' if proc.poll() is None else f'已退出 (退出码: {proc.returncode})'}")
        stop_server(proc)
        return None


def stop_server(proc):
    """停止服务器进程"""
    if proc is None:
        return
    try:
        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        proc.wait(timeout=5)
        print(f"✓ 服务器已停止 (PID: {proc.pid})")
    except ProcessLookupError:
        print(f"服务器进程已不存在 (PID: {proc.pid})")
    except subprocess.TimeoutExpired:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            proc.wait()
            print(f"✓ 服务器已强制停止 (PID: {proc.pid})")
        except Exception as e:
            print(f"停止服务器时出错: {e}")
    except Exception as e:
        print(f"停止服务器时出错: {e}")


def debug_learn_mode(auto_start_server=False):
    """调试模式：学习 tls12_oneway 状态机，遇到目标序列时停止"""
    # 获取模式配置
    mode_config = get_mode_config()
    
    # 初始化日志
    log_dir = "logs"
    os.makedirs(log_dir, exist_ok=True)
    log_file = f"{log_dir}/debug_model_learning_tlse_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    init_global_logger(log_file=log_file, level=logging.INFO, enable_console=True)
    print(f"日志文件: {log_file}")
    
    server_proc = None
    
    try:
        # 自动启动服务器
        if auto_start_server:
            server_proc = start_server(mode_config['server_cmd'])
            if server_proc is None:
                return False
            time.sleep(1)
        
        print("\n" + "=" * 80)
        print("调试模式：TLS 1.2 单向认证状态机学习")
        print(f"模式: {mode_config['description']}")
        print(f"TLS 版本: {mode_config['version']}")
        print(f"需要客户端证书: {mode_config['need_client_cert']}")
        print(f"字母表: {mode_config['alphabet']}")
        print(f"服务器地址: {SERVER_HOST}:{SERVER_PORT}")
        print(f"目标查询序列（遇到时停止）: {TARGET_QUERY}")
        print("=" * 80)
        
        # 创建 TLS 协议配置
        tlsimp = create_tls_protocol(mode_config)
        
        # 创建 SUT
        sul = create_sut(tlsimp, mode_config)
        
        # 确保数据库目录存在
        db_dir = "/tmp/db"
        os.makedirs(db_dir, exist_ok=True)
        db_path = f"{db_dir}/tlse_tls12_oneway_debug.db"
        
        # 创建包装类（用于 L* 算法）
        sul = FailSafeCacheSUL(sul, database=db_path)
        
        # 创建调试包装类（用于检测目标序列）
        debug_sul = DebugStopSUL(sul, TARGET_QUERY)
        
        # 获取字母表
        alphabet = mode_config['alphabet']
        
        print(f"\n开始学习状态机...")
        print(f"字母表大小: {len(alphabet)}")
        print("-" * 80)
        
        # 获取并执行 happy flow 作为初始查询
        happy_flow = get_happy_flow()
        initial_samples = None
        if happy_flow:
            print(f"\n执行 Happy Flow 作为初始查询: {' -> '.join(happy_flow)}")
            try:
                happy_flow_output = debug_sul.query(happy_flow)
                print(f"Happy Flow 响应: {happy_flow_output}")
                initial_samples = [(happy_flow, happy_flow_output)]
                print(f"已将 Happy Flow 添加到初始样本中")
                print("-" * 80)
            except TargetQueryReached as e:
                print(f"\n⚠️  在 Happy Flow 执行过程中遇到目标序列！")
                print(f"目标序列: {e.target_query}")
                print(f"查询数: {e.query_count}")
                if e.output:
                    print(f"响应: {e.output}")
                return True
            except Exception as e:
                print(f"警告: Happy Flow 查询失败: {e}")
                print("继续学习过程（不使用初始样本）...")
        
        # 创建等价性预言
        eq_oracle = StatePrefixOracleFailSafe(alphabet, debug_sul, walks_per_state=1, walk_len=3, database=db_path)
        eq_oracle.add_manual_counterexample(happy_flow)
        
        # 生成输出目录
        imp = tlsimp.implementation
        ver = tlsimp.imp_ver if hasattr(tlsimp, 'imp_ver') else 'None'
        tv = 'tls1.2'
        auth_type = 'oneway'
        folder_name = f"{imp}_{ver}_{tv}_{auth_type}_debug"
        output_base_dir = "dots/tlse"
        os.makedirs(output_base_dir, exist_ok=True)
        folder_path = f"{output_base_dir}/{folder_name}"
        os.makedirs(folder_path, exist_ok=True)
        save_dir = f"{folder_path}/model.dot"
        
        print(f"输出目录: {folder_path}")
        print(f"模型文件: {save_dir}")
        print("-" * 80)
        
        # 运行 L* 算法学习状态机
        print("\n运行 L* 算法...")
        print(f"监控目标序列: {TARGET_QUERY}")
        print("（遇到此序列时将停止学习）")
        print("-" * 80)
        
        try:
            model = run_Lstar(
                alphabet, 
                debug_sul, 
                eq_oracle, 
                'mealy', 
                samples=initial_samples,
                cache_and_non_det_check=False
            )
            
            # 保存模型
            model.save(file_path=save_dir)
            print(f"\n✓ 模型已保存到: {save_dir}")
            
            # 可视化状态机
            print("\n生成状态机可视化...")
            visualize_automaton(model)
            
            print("\n" + "=" * 80)
            print(f"✓ {mode_config['description']} 状态机学习完成！")
            print(f"总查询数: {debug_sul.query_count}")
            print("=" * 80)
            
        except TargetQueryReached as e:
            print("\n" + "=" * 80)
            print("⚠️  学习过程已停止（遇到目标查询序列）")
            print(f"目标序列: {e.target_query}")
            print(f"总查询数: {e.query_count}")
            if e.output:
                print(f"查询响应: {e.output}")
            print("=" * 80)
            
            # 尝试保存当前状态（如果有部分模型）
            print("\n尝试保存当前状态...")
            try:
                # 这里可以添加保存当前观察表或部分模型的逻辑
                print("（当前版本不支持部分模型保存）")
            except Exception as save_error:
                print(f"保存状态时出错: {save_error}")
            
            return True
        
        return True
        
    except Exception as e:
        print(f"\n✗ 学习过程中出错: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        # 停止服务器
        if auto_start_server and server_proc:
            stop_server(server_proc)
            time.sleep(1)


def main():
    parser = argparse.ArgumentParser(
        description='调试模式：TLSE 服务器 TLS 1.2 单向认证状态机学习（遇到目标序列时停止）',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 手动启动服务器后运行调试
  python3 debug_model_learning_tlse.py
  
  # 自动启动服务器并运行调试
  python3 debug_model_learning_tlse.py --auto-start

目标查询序列:
  ('ClientHello', 'ClientKeyExchange', 'ChangeCipherSpec', 'ClosureAlert', 'Finish')
        """
    )
    
    parser.add_argument(
        '--auto-start',
        action='store_true',
        help='自动启动和停止服务器'
    )
    
    args = parser.parse_args()
    
    # 运行调试模式
    debug_learn_mode(auto_start_server=args.auto_start)


if __name__ == '__main__':
    main()

